# -*- coding: utf-8 -*-#

# -------------------------------------------------------------------------------
# Name:         zx-losses
# Description:
# Author:       zx
# Date:         2020/9/30
# -------------------------------------------------------------------------------
import tensorflow as tf


def bce_loss(args, negative_ratio=3.0, eps=1e-6, scale=5.):
    '''
    Args:
        pred: shape :math:`(N, H, W, 1)`, the prediction of network
        gt: shape :math:`(N, H, W)`, the target
        mask: shape :math:`(N, H, W)`, the mask indicates positive regions
    '''
    pred, gt, mask = args
    pred = pred[..., 0]
    positive_mask = (gt * mask)
    negative_mask = ((1 - gt) * mask)
    positive_count = tf.reduce_sum(positive_mask)
    negative_count = tf.reduce_min([tf.reduce_sum(negative_mask), positive_count * negative_ratio])
    loss = tf.keras.backend.binary_crossentropy(gt, pred)
    positive_loss = loss * positive_mask
    negative_loss = loss * negative_mask
    negative_loss, _ = tf.nn.top_k(tf.reshape(negative_loss, (-1,)), tf.cast(negative_count, tf.int32))

    balanced_loss = (tf.reduce_sum(positive_loss) + tf.reduce_sum(negative_loss)) / (
            positive_count + negative_count + eps)
    balanced_loss = balanced_loss * scale
    return balanced_loss


def l1_loss(args, scale=10.):
    pred, gt, mask = args
    pred = pred[..., 0]
    mask_sum = tf.reduce_sum(mask)
    loss = tf.keras.backend.switch(mask_sum > 0, tf.reduce_sum(tf.abs(pred - gt) * mask) / mask_sum, tf.constant(0.))
    loss = loss * scale
    return loss


def dice_loss(args, eps=1e-6):
    pred, gt, mask = args
    pred = pred[..., 0]
    intersection = tf.reduce_sum(pred * gt * mask)
    union = tf.reduce_sum(pred * mask) + tf.reduce_sum(gt * mask) + eps
    loss = 1 - 2.0 * intersection / union
    return loss


def db_loss(args):
    binary, thresh_binary, gt, mask, thresh, thresh_map, thresh_mask = args
    balanced_ce_loss_ = bce_loss([binary, gt, mask])
    l1_loss_ = l1_loss([thresh, thresh_map, thresh_mask])
    dice_loss_ = dice_loss([thresh_binary, gt, mask])
    return l1_loss_ + balanced_ce_loss_ + dice_loss_
